import sys

import baostock as bs
import pandas as pd
import os
from datetime import datetime, timedelta
from baostock.security.sectorinfo import query_stock_concept

def get_hs300_stocks(end_date=None):
    # 登录baostock
    lg = bs.login()
    
    # 获取沪深300成分股
    rs = bs.query_hs300_stocks()
    hs300_stocks = []
    while (rs.error_code == '0') & rs.next():
        hs300_stocks.append(rs.get_row_data())
    
    # 转换为DataFrame
    hs300_df = pd.DataFrame(hs300_stocks, columns=rs.fields)
    
    # 设置日期范围
    end_date = end_date if end_date else datetime.now().strftime('%Y-%m-%d')
    start_date = (datetime.strptime(end_date, '%Y-%m-%d') - timedelta(days=365)).strftime('%Y-%m-%d')
    
    # 创建输出目录
    output_dir = 'hs300_data'
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 创建百分位数据目录
    percentile_dir = os.path.join(output_dir, 'percentile_data')
    if not os.path.exists(percentile_dir):
        os.makedirs(percentile_dir)
    
    # 用于存储所有股票的当日百分位数据
    all_percentile_data = []
    
    # 遍历每只股票获取数据
    for index, row in hs300_df.iterrows():
        code = row['code']
        name = row['code_name']
        
        try:
            # 获取股价信息
            rs_price = bs.query_history_k_data_plus(
                code,
                "date,code,open,high,low,close,volume,amount",
                start_date=start_date,
                end_date=end_date,
                frequency="d",
                adjustflag="2"  # 前复权
            )
            price_df = rs_price.get_data()
            
            # 获取估值信息
            rs_valuation = bs.query_history_k_data_plus(
                code,
                "date,code,peTTM,pbMRQ,psTTM,pcfNcfTTM",
                start_date=start_date,
                end_date=end_date,
                frequency="d",
                adjustflag="2"
            )
            valuation_df = rs_valuation.get_data()
            
            # 获取行业分类信息
            rs_industry = bs.query_stock_industry(code=code)
            industry_info = rs_industry.get_row_data() if rs_industry.error_code == '0' else ['','','','']
            
            # 合并数据
            merged_df = pd.merge(price_df, valuation_df, on=['date', 'code'], how='outer')
            
            # 按日期倒序排序
            merged_df = merged_df.sort_values('date', ascending=False)
            
            # 计算peTTM百分位
            merged_df['peTTM'] = pd.to_numeric(merged_df['peTTM'])
            merged_df['peTTM_percentile'] = merged_df['peTTM'].rank(pct=True)
            
            # 保存完整数据为Excel
            file_path = os.path.join(output_dir, f"{code}_{name}.xlsx")
            merged_df.to_excel(file_path, sheet_name=code,index=False)
            
            # 获取当日数据
            today_data = merged_df[merged_df['date'] == end_date]
            if not today_data.empty:
                # 添加股票名称、行业和概念信息
                today_data['name'] = name
                today_data['industry'] = industry_info[3] if len(industry_info) > 3 else ''
                # 添加到汇总列表
                all_percentile_data.append(today_data)
                print(f"已保存 {code} {name} 当日百分位数据")

            print(f"已保存 {code} {name} 完整数据")
            
        except Exception as e:
            print(f"获取 {code} {name} 数据失败: {str(e)}")
    
    # 将所有股票的百分位数据合并
    if all_percentile_data:
        all_percentile_df = pd.concat(all_percentile_data)
        # 按peTTM百分位排序
        all_percentile_df = all_percentile_df.sort_values('peTTM_percentile', ascending=True)
        # 保存汇总文件
        summary_file = os.path.join(percentile_dir, f"hs300_stocks_percentile_{end_date}.xlsx")
        all_percentile_df.to_excel(summary_file, sheet_name='hs300', index=False)
        print(f"已保存所有股票当日百分位汇总数据到 {summary_file}")
    
    # 登出baostock
    bs.logout()
def get_zz500_stocks(end_date=None):
    """
    获取中证500成分股数据
    """
    # 登录baostock
    bs.login()
    
    # 获取中证500成分股
    rs = bs.query_zz500_stocks()
    zz500_stocks = []
    while (rs.error_code == '0') & rs.next():
        zz500_stocks.append(rs.get_row_data())
    
    # 创建输出目录
    output_dir = "zz500_data"
    os.makedirs(output_dir, exist_ok=True)

    # 创建百分位数据目录
    percentile_dir = os.path.join(output_dir, 'percentile_data')
    if not os.path.exists(percentile_dir):
        os.makedirs(percentile_dir)

    # 设置日期范围
    end_date = end_date if end_date else datetime.now().strftime('%Y-%m-%d')
    start_date = (datetime.now() - timedelta(days=365)).strftime('%Y-%m-%d')
    
    all_percentile_data = []
    
    # 遍历中证500成分股
    for stock in zz500_stocks:
        code = stock[1]
        name = stock[2]
        try:
            # 获取价格数据
            rs_price = bs.query_history_k_data_plus(
                code,
                "date,code,open,high,low,close,preclose,volume,amount",
                start_date=start_date,
                end_date=end_date,
                frequency="d",
                adjustflag="2"
            )
            price_df = rs_price.get_data()
            
            # 获取估值数据
            rs_valuation = bs.query_history_k_data_plus(
                code,
                "date,code,peTTM,pbMRQ,psTTM,pcfNcfTTM",
                start_date=start_date,
                end_date=end_date,
                frequency="d",
                adjustflag="2"
            )
            valuation_df = rs_valuation.get_data()
            
            # 获取行业分类信息
            rs_industry = bs.query_stock_industry(code=code)
            industry_info = rs_industry.get_row_data() if rs_industry.error_code == '0' else ['','','','']
            
            # 合并数据
            merged_df = pd.merge(price_df, valuation_df, on=['date', 'code'], how='outer')
            
            # 按日期倒序排序
            merged_df = merged_df.sort_values('date', ascending=False)
            
            # 计算peTTM百分位
            merged_df['peTTM'] = pd.to_numeric(merged_df['peTTM'])
            merged_df['peTTM_percentile'] = merged_df['peTTM'].rank(pct=True)
            
            # 保存完整数据为Excel
            file_path = os.path.join(output_dir, f"{code}_{name}.xlsx")
            merged_df.to_excel(file_path, sheet_name=code, index=False)
            
            # 获取当日数据
            today_data = merged_df[merged_df['date'] == end_date]
            if not today_data.empty:
                # 添加股票名称和行业信息
                today_data['name'] = name
                today_data['industry'] = industry_info[3] if len(industry_info) > 3 else ''
                # 添加到汇总列表
                all_percentile_data.append(today_data)
                print(f"已保存 {code} {name} 当日百分位数据")
            
            print(f"已保存 {code} {name} 完整数据")
            
        except Exception as e:
            print(f"获取 {code} {name} 数据失败: {str(e)}")
    
    # 将所有股票的百分位数据合并
    if all_percentile_data:
        all_percentile_df = pd.concat(all_percentile_data)
        # 按peTTM百分位排序
        all_percentile_df = all_percentile_df.sort_values('peTTM_percentile', ascending=True)
        # 保存汇总文件
        summary_file = os.path.join(percentile_dir, f"zz500_stocks_percentile_{end_date}.xlsx")
        all_percentile_df.to_excel(summary_file, sheet_name='zz500', index=False)
        print(f"已保存所有股票当日百分位汇总数据到 {summary_file}")
    
    # 登出baostock
    bs.logout()

if __name__ == "__main__":
    day = datetime.now().strftime('%Y-%m-%d')
    if len(sys.argv)==2:
        day = sys.argv[1]
        print(day)
    get_hs300_stocks(end_date=day)
    get_zz500_stocks(end_date=day)
